#  -*- coding: utf-8 -*-

from datetime import datetime

from pandas import DataFrame
from pymongo import ASCENDING
import time
from util.database import DB_CONN
from trading.signal.computer.daily_k_break_ma10_signal_computer import DailyKBreakMA10SignalComputer
from trading.signal.computer.boll_signal_computer import BollSignalComputer


"""
信号管理模块
"""


class SignalModule:
    def get_single_stock_signals(self, code, signal,is_index, begin_date, end_date):
        """
        获取某只股票的某个信号在一段时间内的值
        :param code:  股票代码
        :param signal:  信号名称
        :param begin_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame(columns=['code',signal, 'date'])
        """

        signal_collection = DB_CONN[signal]

        signal_cursor = signal_collection.find(
            {'code': code, 'date': {'$gte': begin_date, '$lte': end_date},"index":is_index},
            sort=[('date', ASCENDING)])

        signal_df = DataFrame(
            [{
                'date': x['date'],
                signal: x[signal],
                'code': x['code'],
                'index': x['index']
            } for x in signal_cursor])

        return signal_df

    def get_single_date_signals(self, signal, date):
        """
        获取某个信号在某个交易日的所有股票的数据
        :param factor: 信号名称
        :param date: 日期
        :return: DataFrame(columns=['code',signal, 'date'])
        """
        signal_collection = DB_CONN[signal]

        signal_cursor = signal_collection.find({'date': date})

        signal_df = DataFrame(
            [{
                'date': x['date'],
                signal: x[signal],
                'code': x['code'],
                'index': x['index']
            } for x in signal_cursor])

        return signal_df

    def compute(self,begin_date=None,end_date=None):
        # 所有的信号实例
        signals = [
            DailyKBreakMA10SignalComputer(),
            BollSignalComputer()
        ]

        now = datetime.now().strftime('%Y-%m-%d')
        if begin_date is None:
            begin_date = "1995-02-24"
        if end_date is None:
            end_date = now
        for signal in signals:
            start_time = time.time()

            signal.compute(begin_date=begin_date, end_date=end_date)
            end_time = time.time()
            print('计算信号：%s, 从日期：%s 到 %s  ，共耗时：%.3f 秒' % (signal.name, begin_date, end_date, (end_time - start_time)), flush=True)
